## 从两个文件夹中每隔 n 张取一张图片，放到新的文件夹中。
## 按顺序重命名图片文件和标签文件
import os
import shutil
import xml.etree.ElementTree as ET

def select_images(folder1, folder2, output_folder1, output_folder2, n):
    """
    从两个文件夹中每隔 n 张取一张图片，放到新的文件夹中。
    
    Args:
        folder1: 第一个文件夹路径
        folder2: 第二个文件夹路径
        output_folder1: 新建的第一个输出文件夹路径
        output_folder2: 新建的第二个输出文件夹路径
        n: 每隔 n 张取一张图片
    """
    # 确保输出文件夹存在
    os.makedirs(output_folder1, exist_ok=True)
    os.makedirs(output_folder2, exist_ok=True)

    # 获取两个文件夹中的图片文件列表，并按名称排序（确保对应关系）
    files1 = sorted([f for f in os.listdir(folder1) if os.path.isfile(os.path.join(folder1, f))])
    files2 = sorted([f for f in os.listdir(folder2) if os.path.isfile(os.path.join(folder2, f))])

    # 检查文件数量是否相等
    if len(files1) != len(files2):
        print("两个文件夹中的文件数量不匹配，请检查数据集！")
        return

    # 按间隔 n 取文件
    for i in range(0, len(files1), n):
        # 获取当前文件路径
        file1_path = os.path.join(folder1, files1[i])
        file2_path = os.path.join(folder2, files2[i])

        # 复制文件到新文件夹
        shutil.copy(file1_path, os.path.join(output_folder1, files1[i]))
        shutil.copy(file2_path, os.path.join(output_folder2, files2[i]))

    print(f"操作完成！已将每隔 {n} 张图片保存到 {output_folder1} 和 {output_folder2}。")

def rename_images_and_update_annotations(images_dir, annotations_dir, start_num):
    """
    重命名图片文件并更新对应的标签文件。

    :param images_dir: 图片文件夹路径
    :param annotations_dir: 标签文件夹路径
    """
    # 获取图片文件名，按顺序排序
    image_files = sorted([f for f in os.listdir(images_dir) if f.endswith('.jpg')])
    annotation_files = sorted([f for f in os.listdir(annotations_dir) if f.endswith('.xml')])

    # 确保图片和标签数量一致
    if len(image_files) != len(annotation_files):
        raise ValueError("图片和标签文件数量不一致！")

    # 遍历图片文件，按顺序重命名图片文件，并修改对应的标签文件
    for i, (image_file, annotation_file) in enumerate(zip(image_files, annotation_files)):
        # 新的图片名称
        new_image_name = f"frame_{(i + start_num):05d}.jpg"
        new_annotation_name = f"frame_{(i + start_num):05d}.xml"
        
        # 重命名图片
        old_image_path = os.path.join(images_dir, image_file)
        new_image_path = os.path.join(images_dir, new_image_name)
        os.rename(old_image_path, new_image_path)

        # 重命名标签文件
        old_annotation_path = os.path.join(annotations_dir, annotation_file)
        new_annotation_path = os.path.join(annotations_dir, new_annotation_name)
        os.rename(old_annotation_path, new_annotation_path)

        # 处理对应的XML标签文件
        tree = ET.parse(new_annotation_path)
        root = tree.getroot()

        # 修改<filename>标签
        filename_tag = root.find('filename')
        if filename_tag is not None:
            filename_tag.text = new_image_name
        
        # 删除<path>标签（如果存在）
        path_tag = root.find('path')
        if path_tag is not None:
            root.remove(path_tag)

        # 保存修改后的XML文件
        tree.write(new_annotation_path)

    print("图片和标签文件重命名完成！")

# 示例使用
folder1 = "/home/shuai/Downloads/total_202411272006/images"  # 替换为第一个文件夹的路径
folder2 = "/home/shuai/Downloads/total_202411272006/annotations"  # 替换为第二个文件夹的路径
output_folder1 = "/home/shuai/Downloads/total_202411272006/select_image/paper_image_label/images"  # 替换为第一个输出文件夹的路径
output_folder2 = "/home/shuai/Downloads/total_202411272006/select_image/paper_image_label/annotations"  # 替换为第二个输出文件夹的路径

## 从两个文件夹中每隔 n 张取一张图片，放到新的文件夹中。
# n = 10  # 每隔 n 张取一张
# select_images(folder1, folder2, output_folder1, output_folder2, n)

## 按顺序重命名图片文件和标签文件
rename_images_and_update_annotations(output_folder1, output_folder2, start_num = 0)
